#!/usr/bin/env python
# -*- coding: utf-8 -*-

import asyncio
import logging
from typing import List, Dict, Any, Optional
import numpy as np
import aiohttp
import json

try:
    import dashscope
    from dashscope import TextEmbedding
except ImportError:
    dashscope = None
    TextEmbedding = None

from app.core.config import settings
from app.core.exceptions import EmbeddingError


logger = logging.getLogger(__name__)


class EmbeddingManager:
    """多提供商向量化管理器
    
    支持多种向量化服务提供商，包括阿里云和Ollama。
    """
    
    def __init__(self):
        """初始化向量化管理器"""
        self.provider = settings.embedding_provider
        self.dimension = settings.embedding_dimension  # 从配置读取向量维度
        
        # 阿里云配置
        self.alibaba_model = settings.alibaba_embedding_model
        self.alibaba_api_key = settings.alibaba_api_key
        
        # Ollama配置
        self.ollama_base_url = settings.ollama_base_url
        self.ollama_model = settings.ollama_embedding_model
        
        # 初始化对应的提供商
        if self.provider == "alibaba":
            self._init_alibaba()
        elif self.provider == "ollama":
            self._init_ollama()
        else:
            raise EmbeddingError(f"不支持的向量化提供商: {self.provider}")
        
        logger.info(f"已初始化向量化管理器，提供商: {self.provider}")
    
    def _init_alibaba(self):
        """初始化阿里云向量化服务"""
        if not self.alibaba_api_key:
            logger.warning("未配置阿里云API密钥，向量化功能将不可用")
            return
        
        if not dashscope:
            logger.error("未安装dashscope包，请运行: pip install dashscope")
            raise EmbeddingError("缺少dashscope依赖")
        
        # 设置API密钥
        dashscope.api_key = self.alibaba_api_key
        logger.info(f"已初始化阿里云向量化服务，模型: {self.alibaba_model}")
    
    def _init_ollama(self):
        """初始化Ollama向量化服务"""
        if not self.ollama_base_url:
            raise EmbeddingError("未配置Ollama服务地址")
        
        logger.info(f"已初始化Ollama向量化服务，地址: {self.ollama_base_url}, 模型: {self.ollama_model}")
    
    async def embed_text(self, text: str) -> List[float]:
        """
        对单个文本进行向量化
        
        Args:
            text: 要向量化的文本
            
        Returns:
            List[float]: 文本的向量表示
            
        Raises:
            EmbeddingError: 当向量化失败时
        """
        if not text.strip():
            logger.warning("输入文本为空，返回零向量")
            return [0.0] * self.dimension
        
        if self.provider == "alibaba":
            return await self._embed_text_alibaba(text)
        elif self.provider == "ollama":
            return await self._embed_text_ollama(text)
        else:
            raise EmbeddingError(f"不支持的向量化提供商: {self.provider}")
    
    async def _embed_text_alibaba(self, text: str) -> List[float]:
        """使用阿里云进行文本向量化"""
        if not self.alibaba_api_key or not dashscope:
            raise EmbeddingError("阿里云向量化服务未正确配置")
        
        try:
            # 使用asyncio.to_thread在线程池中运行同步函数
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: TextEmbedding.call(
                    model=self.alibaba_model,
                    input=text
                )
            )
            
            if response.status_code != 200:
                error_msg = f"阿里云向量化请求失败: {response.status_code}, {response.message}"
                logger.error(error_msg)
                raise EmbeddingError(error_msg)
            
            if not response.output or not response.output.get('embeddings'):
                raise EmbeddingError("阿里云向量化响应格式错误")
            
            embeddings = response.output['embeddings']
            if not embeddings:
                raise EmbeddingError("未获取到向量数据")
            
            # 获取第一个文本的向量
            vector = embeddings[0]['embedding']
            
            logger.debug(f"阿里云文本向量化成功，维度: {len(vector)}")
            return vector
            
        except Exception as e:
            error_msg = f"阿里云文本向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def _embed_text_ollama(self, text: str) -> List[float]:
        """使用Ollama进行文本向量化"""
        try:
            url = f"{self.ollama_base_url}/api/embed"
            payload = {
                "model": self.ollama_model,
                "input": [text]
            }
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as response:
                    if response.status != 200:
                        error_msg = f"Ollama向量化请求失败: {response.status}"
                        logger.error(error_msg)
                        raise EmbeddingError(error_msg)
                    
                    result = await response.json()
                    
                    if not result.get('embeddings'):
                        raise EmbeddingError("Ollama向量化响应格式错误")
                    
                    embeddings = result['embeddings']
                    if not embeddings:
                        raise EmbeddingError("未获取到向量数据")
                    
                    # 获取第一个文本的向量
                    vector = embeddings[0]
                    
                    logger.debug(f"Ollama文本向量化成功，维度: {len(vector)}")
                    return vector
                    
        except Exception as e:
            error_msg = f"Ollama文本向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """
        对多个文本进行批量向量化
        
        Args:
            texts: 要向量化的文本列表
            
        Returns:
            List[List[float]]: 文本向量列表
            
        Raises:
            EmbeddingError: 当向量化失败时
        """
        if not texts:
            return []
        
        # 过滤空文本
        non_empty_texts = [text for text in texts if text.strip()]
        if not non_empty_texts:
            logger.warning("所有输入文本都为空，返回零向量列表")
            return [[0.0] * self.dimension] * len(texts)
        
        if self.provider == "alibaba":
            return await self._embed_texts_alibaba(texts)
        elif self.provider == "ollama":
            return await self._embed_texts_ollama(texts)
        else:
            raise EmbeddingError(f"不支持的向量化提供商: {self.provider}")
    
    async def _embed_texts_alibaba(self, texts: List[str]) -> List[List[float]]:
        """使用阿里云进行批量文本向量化"""
        if not self.alibaba_api_key or not dashscope:
            raise EmbeddingError("阿里云向量化服务未正确配置")
        
        # 过滤空文本
        non_empty_texts = [text for text in texts if text.strip()]
        if not non_empty_texts:
            return [[0.0] * self.dimension] * len(texts)
        
        try:
            # 批量处理，考虑API限制
            batch_size = 25  # 阿里云API的批量限制
            all_vectors = []
            
            for i in range(0, len(non_empty_texts), batch_size):
                batch_texts = non_empty_texts[i:i + batch_size]
                
                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(
                    None,
                    lambda: TextEmbedding.call(
                        model=self.alibaba_model,
                        input=batch_texts
                    )
                )
                
                if response.status_code != 200:
                    error_msg = f"阿里云批量向量化请求失败: {response.status_code}, {response.message}"
                    logger.error(error_msg)
                    raise EmbeddingError(error_msg)
                
                if not response.output or not response.output.get('embeddings'):
                    raise EmbeddingError("阿里云批量向量化响应格式错误")
                
                embeddings = response.output['embeddings']
                batch_vectors = [emb['embedding'] for emb in embeddings]
                all_vectors.extend(batch_vectors)
                
                # 避免API频率限制
                if i + batch_size < len(non_empty_texts):
                    await asyncio.sleep(0.1)
            
            # 处理原始文本中的空文本
            result_vectors = []
            vector_index = 0
            
            for text in texts:
                if text.strip():
                    result_vectors.append(all_vectors[vector_index])
                    vector_index += 1
                else:
                    result_vectors.append([0.0] * self.dimension)
            
            logger.info(f"阿里云批量向量化成功，处理了{len(texts)}个文本")
            return result_vectors
            
        except Exception as e:
            error_msg = f"阿里云批量向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def _embed_texts_ollama(self, texts: List[str]) -> List[List[float]]:
        """使用Ollama进行批量文本向量化"""
        # 过滤空文本
        non_empty_texts = [text for text in texts if text.strip()]
        if not non_empty_texts:
            return [[0.0] * self.dimension] * len(texts)
        
        try:
            # Ollama可以一次处理多个文本
            url = f"{self.ollama_base_url}/api/embed"
            payload = {
                "model": self.ollama_model,
                "input": non_empty_texts
            }
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as response:
                    if response.status != 200:
                        error_msg = f"Ollama批量向量化请求失败: {response.status}"
                        logger.error(error_msg)
                        raise EmbeddingError(error_msg)
                    
                    result = await response.json()
                    
                    if not result.get('embeddings'):
                        raise EmbeddingError("Ollama批量向量化响应格式错误")
                    
                    all_vectors = result['embeddings']
                    
                    # 处理原始文本中的空文本
                    result_vectors = []
                    vector_index = 0
                    
                    for text in texts:
                        if text.strip():
                            result_vectors.append(all_vectors[vector_index])
                            vector_index += 1
                        else:
                            result_vectors.append([0.0] * self.dimension)
                    
                    logger.info(f"Ollama批量向量化成功，处理了{len(texts)}个文本")
                    return result_vectors
                    
        except Exception as e:
            error_msg = f"Ollama批量向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def embed_query(self, query: str) -> List[float]:
        """
        对查询文本进行向量化（可能需要特殊处理）
        
        Args:
            query: 查询文本
            
        Returns:
            List[float]: 查询向量
        """
        # 对查询可能需要特殊的前缀处理
        processed_query = f"查询: {query}" if len(query) < 100 else query
        return await self.embed_text(processed_query)
    
    def calculate_similarity(self, vector1: List[float], vector2: List[float]) -> float:
        """
        计算两个向量的余弦相似度
        
        Args:
            vector1: 第一个向量
            vector2: 第二个向量
            
        Returns:
            float: 余弦相似度（-1到1之间）
        """
        try:
            v1 = np.array(vector1)
            v2 = np.array(vector2)
            
            # 计算余弦相似度
            dot_product = np.dot(v1, v2)
            norm1 = np.linalg.norm(v1)
            norm2 = np.linalg.norm(v2)
            
            if norm1 == 0 or norm2 == 0:
                return 0.0
            
            similarity = dot_product / (norm1 * norm2)
            return float(similarity)
            
        except Exception as e:
            logger.error(f"计算相似度失败: {str(e)}")
            return 0.0
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        if self.provider == "alibaba":
            return {
                "model_name": self.alibaba_model,
                "provider": "alibaba",
                "dimension": self.dimension,
                "max_batch_size": 25,
                "available": self.alibaba_api_key is not None and dashscope is not None
            }
        elif self.provider == "ollama":
            return {
                "model_name": self.ollama_model,
                "provider": "ollama",
                "dimension": self.dimension,
                "max_batch_size": 100,  # Ollama通常支持更大的批量
                "available": self.ollama_base_url is not None
            }
        else:
            return {
                "model_name": "unknown",
                "provider": self.provider,
                "dimension": self.dimension,
                "max_batch_size": 1,
                "available": False
            }
    
    async def test_connection(self) -> bool:
        """
        测试向量化服务连接
        
        Returns:
            bool: 连接是否正常
        """
        try:
            test_text = "测试连接"
            vector = await self.embed_text(test_text)
            return len(vector) > 0
        except Exception as e:
            logger.error(f"向量化服务连接测试失败: {str(e)}")
            return False


# 创建全局向量化管理器实例
embedding_manager = EmbeddingManager()